#!/usr/bin/python
# -*- coding: utf-8 -*-
import os
import argparse
import psycopg2
from dotenv import load_dotenv
from psycopg2 import sql

# 加载环境变量
load_dotenv(".env")


def create_geomtry_table(cursor, table_name):
  """
  创建包含几何字段的表
  :param cursor: 数据库游标
  :tableName: 表名
  :param geomType:几何类型, 如'POINT', 'LINESTRING', 'POLYGON'.
                  格式为字符串，如'GEOMETRY(Point, 4326)'
          id SERIAL PRIMARY KEY,
          name character varying(255) NOT NULL,
          geom {geometry_type}
  """
  # 创建表的SQL语句
  create_table_sql = sql.SQL("""
      CREATE TABLE IF NOT EXISTS {table_name}();
  """).format(
      table_name=sql.Identifier(table_name)
  )
  # 执行SQL语句
  cursor.execute(create_table_sql)


def add_columns_to_table(cur, table_name, columns):
  """
  向指定的表中添加新列。
  :param cur: psycopg2数据库连接游标
  :param tableName: 要添加新列的表名
  :param columns: 包含新列名和数据类型的列表，每个元素是一个(列名, 数据类型)的元组
  """
  # 示例用法（注意：这里没有包括数据库连接的代码）
  # cursor = ...  # 获取游标对象
  # table_name = 'your_table'
  # columns_to_add = [('new_column1', 'INTEGER'), ('new_column2', 'VARCHAR(100)')]
  # add_columns_to_table(cursor, table_name, columns_to_add)
  # 遍历列名列表，为每个新列构建ALTER TABLE语句
  for column_name, data_type in columns:
    add_column_query = sql.SQL("""
      ALTER TABLE {table}
      ADD COLUMN {column} {type}
      """).format(
        table=sql.Identifier(table_name),
        column=sql.Identifier(column_name),
        type=sql.SQL(data_type)
    )
    # 执行ALTER TABLE语句
    cur.execute(add_column_query)


def init_database():
    try:
        # 从环境变量获取配置
        dbname = os.getenv("DB_NAME")
        user = os.getenv("DB_USER")
        password = os.getenv("DB_PASSWORD")
        host = os.getenv("DB_HOST", "localhost")
        port = os.getenv("DB_PORT", 5432)
        table_name=os.getenv("TABLE_NAME")

        # 验证必要参数
        if not all([dbname, user, password, port]):
            raise ValueError("Missing required database configuration in .env file")

        # 连接到默认数据库创建目标数据库
        conn = psycopg2.connect(
            dbname="postgres",
            user=user,
            password=password,
            host=host,
            port=port
        )
        conn.autocommit = True
        cur = conn.cursor()

        # 检查数据库是否已存在
        cur.execute(
        sql.SQL("SELECT 1 FROM pg_database WHERE datname = %s"),
        [dbname])
        exists = cur.fetchone()

        # 创建新数据库
        if not exists:
            cur.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(dbname)))
            print(f"Database '{dbname}' created successfully.")
        else:
            print(f"Database '{dbname}' already exists.")

        # 重新连接数据库
        cur.close()
        conn.close()
        conn = psycopg2.connect(
            dbname=dbname,
            user=user,
            password=password,
            host=host,
            port=port
        )
        conn.autocommit = True
        cur = conn.cursor()

        # 创建PostGIS扩展
        print("Create PostGIS extension...")
        cur.execute("CREATE EXTENSION IF NOT EXISTS postgis")

        # 创建数据表
        print("Create table...")
        create_geomtry_table(cur, table_name)
        # 向表中添加列
        columns_to_add = [('name', 'VARCHAR(50)'), ('type', 'CHAR(10)'),
                          ('x', 'DOUBLE PRECISION'), ('y', 'DOUBLE PRECISION'),
                          ('depth', 'REAL'),
                          ('ground_elevation', 'DOUBLE PRECISION'), ('bottom_elevation', 'DOUBLE PRECISION'),
                          ('cover_shape', 'VARCHAR(20)'), ('cover_size_1', 'REAL'), ('cover_size_2', 'REAL'),
                          ('road_name', 'VARCHAR(50)'),
                          ('geometry', 'GEOMETRY')]
        print("Add columns to table...")
        add_columns_to_table(cur, table_name, columns_to_add)

        # 创建空间索引
        print("Create index...")
        create_spatial_index = sql.SQL("""
            CREATE INDEX geom_idx
            ON {} USING GIST(geometry)
            """).format(sql.Identifier(table_name))
        cur.execute(create_spatial_index)
        print(f"Database {dbname} initialized successfully!")
    except Exception as e:
        print(f"初始化失败: {str(e)}")
        conn.rollback()
        exit(1)

    finally:
        if 'cur' in locals(): cur.close()
        if 'conn' in locals(): conn.close()


if __name__ == "__main__":
    init_database()
